"""
    promptTemplate工厂
"""
from langchain.prompts import (
    PromptTemplate,
    FewShotPromptTemplate,
)


def createFewShotPromptTemplate(examples, example_prompt, prefix, suffix, input_variables):
    """
    创建少样本提示词模板
    :param examples:
    :param example_prompt:
    :param prefix:
    :param suffix:
    :param input_variables:
    :return:
    """
    return FewShotPromptTemplate(
        examples=examples,
        example_prompt=example_prompt,
        example_separator="\n",
        prefix=prefix,
        suffix=suffix,
        input_variables=input_variables,
    )


class PromptTemplateFactory:

    @staticmethod
    def testPromptTemplate(query, examples):
        prefix = "你是个历史问答助手,但是你的回答必须参考下文提供的史料,并且下文提供的史料是按照相关度从高到低排序。如果在史料中没有找到答案，回答不知道。"
        suffix = f"请根据上文的史料回答我的问题：{query}"
        example_separator = "\n",
        input_variables = []
        example_prompt = \
            PromptTemplate(input_variables=["example"],
                           template="史料：{example}"
                           )
        return createFewShotPromptTemplate(examples, example_prompt, prefix, suffix, input_variables)

    @staticmethod
    def boatNormPromptTemplate(query, examples):
        prefix = "你是个船舶行业的知识问答助手,但是你的回答必须参考下文提供的法规,并且下文提供的法规是按照相关度从高到低排序。如果在法规中没有找到答案，回答不知道。"
        suffix = f"请根据上文的法规回答我的问题：{query}"
        example_separator = "\n",
        input_variables = []
        example_prompt = \
            PromptTemplate(input_variables=["example"],
                           template="法规：{example}"
                           )
        return createFewShotPromptTemplate(examples, example_prompt, prefix, suffix, input_variables)


if __name__ == '__main__':
    from database_helper import milvusDatabase
    # 问题向量化
    query = "历史上被称作贼将军的是谁？"

    from llm_chat_client import qianfanLLM
    embeddings = qianfanLLM.embed_query(query)

    from param_builder import MilvusParam
    db_param = MilvusParam.getUserParam()
    ans = milvusDatabase.searchEmbedding([embeddings], db_param)
    examples = []
    for item in ans[0]:
        examples.append({"example": item.fields["desc"]})

    prompt = PromptTemplateFactory.testPromptTemplate(query, examples)
    print(prompt.format_prompt())

    result = qianfanLLM.invoke(prompt)
    print(type(result.content))
    print(result)
